from matplotlib import pyplot as plt


class TimeDrainPlot:
    def __init__(self):
        self.originDataSet = []
        self.alignDataSet = []

    def parseFile(self, model, type):
        if(type == "origin"):
            self.originDataSet.append(model.timeAndDrain)
        else:
            self.alignDataSet.append(model.timeAndDrain)

    def plot(self):
        ax = plt.figure().add_subplot(111)
        index = 1
        for data in self.originDataSet:
            y_start = data["drain"][0]
            x_start = data["time"][0]
            x_list = [item - x_start for item in data["time"]]
            y_list = [item - y_start for item in data["drain"]]
            print(y_list)
            ax.plot(x_list,y_list ,label="Origin" + str(index), color = 'red')
            index += 1
        index = 1
        for data in self.alignDataSet:
            y_start = data["drain"][0]
            x_start = data["time"][0]
            x_list = [item - x_start for item in data["time"]]
            y_list = [item - y_start for item in data["drain"]]
            print(y_list)
            ax.plot(x_list, y_list, label="Align" + str(index), color='deepskyblue')
            index += 1
        ax.set_xlabel("Time", fontproperties='Times New Roman', fontsize=12)
        ax.set_ylabel("Drain(mAh)", fontproperties='Times New Roman', fontsize=12)
        ax.legend(bbox_to_anchor=(1.05, 0), loc=3, borderaxespad=0)
        plt.show()